#define lib_serialize_c
#define LUA_LIB

#include "lua.h"
#include "lauxlib.h"
#include "lualib.h"

#include "lj_obj.h"
#include "lj_lib.h"
#include "lib_serialize.h"
#include <assert.h>

#define LUA_SERIALIZE_CONTEXT "serialize_context*"

uint32_t const TOP = (1<<24);
uint32_t const BOT = (1<<15);

/* ------------------------------------------------------------------------ */

#define LJLIB_MODULE_serialize_method

static serialize_context *serialize_toctx(lua_State *L)
{
	if (!(L->base < L->top && tvisudata(L->base) &&
		udataV(L->base)->udtype == UDTYPE_SERIALIZE_CONTEXT))
		lj_err_argtype(L, 1, "serialize_context*");
	return (serialize_context *)uddata(udataV(L->base));
}

LJLIB_CF(serialize_method___gc)
{
	serialize_context *ctx = serialize_toctx(L);
	
	bytewriter_destroy(&ctx->writer);
	bytereader_destroy(&ctx->reader);
	return 0;
}

LJLIB_PUSH(top-1) LJLIB_SET(__index)

/* ------------------------------------------------------------------------ */

#define LJLIB_MODULE_serialize

static unsigned char get_reserve(serialize_context* ctx)
{
	if(ctx->reader.next != ctx->reader.end)
		return *ctx->reader.next++;
	ctx->reader.underflow(&ctx->reader);
	if(ctx->reader.next != ctx->reader.end)
	{
		assert(ctx->empty_bytes + 1 <= (ctx->reader.end - ctx->reader.next));
		// TODO: Is it possible that (empty_bytes + 1 > (ctx->reader.end - ctx->reader.next))?
		// We need a sanity check if that's the case.
		while(ctx->empty_bytes > 0)
		{
			// Fill incorrect zero-bytes
			--ctx->empty_bytes;
			ctx->d_code += (*ctx->reader.next++) << (ctx->empty_bytes * 8);
		}

		return *ctx->reader.next++;
	}

	// Return zero-bytes that need to be filled later on
	++ctx->empty_bytes;
	return 0;
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_putend(serialize_context* ctx)
{
	// Increment tail of queue
	++ctx->message_end_queue;
}

static void serialize_reset_d(serialize_context* ctx)
{
	int i;
	ctx->d_code = 0;
	ctx->d_low = 0;
	ctx->d_range = 0xffffffff;
	ctx->empty_bytes = 0;

	// Sections are always at least 4 bytes
	for(i = 0; i < 4; ++i)
	{
		ctx->d_code = (ctx->d_code << 8) + bytereader_get(&ctx->reader);
	}
}

LJ_FUNCA int32_t LJ_FASTCALL lj_serialize_firstmsg(serialize_context* ctx)
{
	ctx->reader.update(&ctx->reader);
	if(ctx->cur_message_end == ctx->reader.next_unreceived_message)
		return 0;

	// Do we need to read in another section?
	if(ctx->cur_message_end == ctx->reader.next_complete_message_end)
	{
		int ret = ctx->reader.underflow(&ctx->reader);
		assert(ret);
		serialize_reset_d(ctx);
	}

	return 1;
}

LJ_FUNCA int32_t LJ_FASTCALL lj_serialize_nextmsg(serialize_context* ctx)
{
	if(++ctx->cur_message_end == ctx->reader.next_unreceived_message)
		return 0;

	if(ctx->cur_message_end == ctx->reader.next_complete_message_end)
	{
		int ret = ctx->reader.underflow(&ctx->reader);
		assert(ret);
		serialize_reset_d(ctx);
	}

	return 1;
}


static void serialize_renorm_c(serialize_context* ctx)
{
	while ((ctx->c_low ^ (ctx->c_low+ctx->c_range)) < TOP || ctx->c_range < BOT &&
        ((ctx->c_range = (0-ctx->c_low) & (BOT-1)),1))
	{
		bytewriter_put(&ctx->writer, ctx->c_low >> 24);

		// Important that incrementing is done after the byte is written.
		// 2^10 message ends per byte is more than enough in practice, but we have to make sure it's theoretically
		// enough.
		ctx->writer.message_ends += (ctx->message_end_queue >> 22);
		ctx->message_end_queue <<= 11;

		ctx->c_range <<= 8;
		ctx->c_low <<= 8;
	}
}

static void serialize_renorm_d(serialize_context* ctx)
{
	while ((ctx->d_low ^ (ctx->d_low+ctx->d_range)) < TOP || ctx->d_range < BOT &&
        ((ctx->d_range = (0-ctx->d_low) & (BOT-1)),1))
	{
		//ctx->d_code = (ctx->d_code << 8) + bytereader_get(&ctx->reader);
		unsigned char n = get_reserve(ctx);
		ctx->d_code = (ctx->d_code << 8) + n;
		ctx->d_range <<= 8;
		ctx->d_low <<= 8;
	}
}

static void serialize_flush(serialize_context* ctx)
{
	uint32_t out, high, diff;
	int i;
	if(!ctx->message_end_queue && !ctx->writer.message_ends)
		return; // Nothing

	// If we didn't get a message end after the last written byte, there's
	// still a message in progress. That could cause complete sections without
	// message ends in the section logic (and wouldn't work anyway).
	assert((ctx->message_end_queue & ((1<<11)-1)) > 0 && "Expected a message end after last written byte");

#if 0
	for(i = 0; i < 3; ++i)
	{
		bytewriter_put(&ctx->writer, ctx->c_low >> 24);
		ctx->writer.message_ends += (ctx->message_end_queue >> 22);
		ctx->message_end_queue <<= 11;

		ctx->c_low <<= 8;
	}
#else
	high = (ctx->c_low + ctx->c_range - 1);
	diff = ctx->c_low ^ high;

	diff |= diff >> 16;
	diff |= diff >> 8;
	diff |= diff >> 4;
	diff |= diff >> 2;
	diff |= diff >> 1;

	// Remove unnecessary bits
	out = high & ~(diff >> 1);

	// TODO: The section logic assumes that this doesn't write more than
	// 3 bytes as that could create a new complete section without message ends.

	while(out)
	{
		bytewriter_put(&ctx->writer, out >> 24);
		ctx->writer.message_ends += (ctx->message_end_queue >> 22);
		ctx->message_end_queue <<= 11;

		out <<= 8;
	}
#endif

	// Make sure all message ends are accounted for, even if we
	// expected them later.
	while(ctx->message_end_queue)
	{
		ctx->writer.message_ends += (ctx->message_end_queue >> 22);
		ctx->message_end_queue <<= 11;
	}

	ctx->c_low = 0;
	ctx->c_range = 0xffffffff;
}

/*
TODO: A few different encoding methods:

* Integer encoding with chosen context. Can be used efficiently for, for example,
  encoding type ids in an action stream.
* General purpose byte encoding. Something akin to sr2, to be used when it's unclear
  how to otherwise handle the data.
*/

LJ_FUNCA void LJ_FASTCALL lj_serialize_putbit(serialize_context* ctx, int32_t y, int32_t p)
{
	uint32_t r = (ctx->c_range >> 12) * p;
	uint32_t mask = y - 1;

	ctx->c_low += r & mask;
	ctx->c_range = (ctx->c_range & mask) + (r ^ mask) - mask;

	serialize_renorm_c(ctx);
}

LJ_FUNCA int32_t LJ_FASTCALL lj_serialize_getbit(serialize_context* ctx, int32_t p)
{
	uint32_t r = (ctx->d_range >> 12) * p;

	int32_t y = (ctx->d_code - ctx->d_low) < r;
	uint32_t mask = y - 1;

	ctx->d_low += r & mask;
	ctx->d_range = (ctx->d_range & mask) + (r ^ mask) - mask;
		
	serialize_renorm_d(ctx);

	return y;
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_putunibit(serialize_context* ctx, int32_t y)
{
	ctx->c_low += (ctx->c_range >>= 1) & (y - 1);

	serialize_renorm_c(ctx);
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_encode_range(serialize_context* ctx, uint32_t base, uint32_t size, uint32_t total)
{
	uint32_t r = (ctx->c_range / total);

	assert(size > 0 && base + size <= total);
	assert(r > 0);
	
	ctx->c_low += r * base;
	ctx->c_range = r * size;

	serialize_renorm_c(ctx);
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_encode_shift(serialize_context* ctx, uint32_t base, uint32_t size, int32_t sh)
{
	uint32_t r = (ctx->c_range >> sh);

	assert(size > 0 && base + size <= (1<<sh));
	assert(r > 0);
	assert(sh <= 15);

	ctx->c_low += r * base;
	ctx->c_range = r * size;

	serialize_renorm_c(ctx);
}

LJ_FUNCA int32_t LJ_FASTCALL lj_serialize_getunibit(serialize_context* ctx)
{
	//uint32_t y = (ctx->d_code - ctx->d_low - (ctx->d_range >>= 1)) >> 31;
	uint32_t y = (ctx->d_code - ctx->d_low) < (ctx->d_range >>= 1);

	ctx->d_low += ctx->d_range & (y - 1);

	serialize_renorm_d(ctx);

	return (int32_t)y;
}

LJ_FUNCA uint32_t LJ_FASTCALL lj_serialize_get_count(serialize_context* ctx, uint32_t total)
{
	return (ctx->d_code - ctx->d_low) / (ctx->d_range /= total);
}

LJ_FUNCA uint32_t LJ_FASTCALL lj_serialize_get_count_shift(serialize_context* ctx, int sh)
{
	return (ctx->d_code - ctx->d_low) / (ctx->d_range >>= sh);
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_decoded_range(serialize_context* ctx, uint32_t base, uint32_t size)
{
	ctx->d_low += ctx->d_range * base;
	ctx->d_range *= size;

	serialize_renorm_d(ctx);
}

LJ_FUNCA void LJ_FASTCALL lj_serialize_putint(serialize_context* ctx, int32_t i)
{
	lj_serialize_encode_shift(ctx, (i >> 21) & ((1<<11)-1), 1, 11);
	lj_serialize_encode_shift(ctx, (i >> 10) & ((1<<11)-1), 1, 11);
	lj_serialize_encode_shift(ctx, i & ((1<<10)-1), 1, 10);
}

LJ_FUNCA int32_t LJ_FASTCALL lj_serialize_getint(serialize_context* ctx)
{
	uint32_t base1, base2, base3;
	base1 = lj_serialize_get_count_shift(ctx, 11);
	lj_serialize_decoded_range(ctx, base1, 1);
	base2 = lj_serialize_get_count_shift(ctx, 11);
	lj_serialize_decoded_range(ctx, base2, 1);
	base3 = lj_serialize_get_count_shift(ctx, 10);
	lj_serialize_decoded_range(ctx, base3, 1);

	return (int32_t)((base1 << 21) + (base2 << 10) + base3);
}

static void serialize_flush_handler(bytewriter* w)
{
	serialize_context* ctx = (serialize_context*)((char*)w - offsetof(serialize_context, writer));

	serialize_flush(ctx);
}

serialize_context* serialize_create(lua_State* L)
{
	serialize_context* ctx = (serialize_context *)lua_newuserdata(L, sizeof(serialize_context));
	GCudata *ud = udataV(L->top-1);
	ud->udtype = UDTYPE_SERIALIZE_CONTEXT;

	lua_getfield(L, LUA_REGISTRYINDEX, LUA_SERIALIZE_CONTEXT); // ctx mt
	lua_setmetatable(L, -2); // ctx
	ctx->status = 0;

	ctx->c_low = 0;
	ctx->c_range = 0xffffffff;

	ctx->d_low = 0;
	ctx->d_code = 0;
	ctx->d_range = 0xffffffff;

	ctx->message_end_queue = 0;
	ctx->cur_message_end = 0;
	ctx->empty_bytes = 0;

	bytewriter_init(&ctx->writer);
	bytereader_init(&ctx->reader);

	ctx->writer.flush = serialize_flush_handler;
	return ctx;
}

LJLIB_CF(serialize_putint)  LJLIB_REC(serialize_putint IRCALL_lj_serialize_putint)
{
	serialize_context* ctx = serialize_toctx(L);
  
	lj_serialize_putint(ctx, lj_lib_checkint(L, 2));
	return 0;
}

LJLIB_CF(serialize_putend)  LJLIB_REC(serialize_putend IRCALL_lj_serialize_putend)
{
	serialize_context* ctx = serialize_toctx(L);
  
	lj_serialize_putend(ctx);
	return 0;
}

LJLIB_CF(serialize_getint)  LJLIB_REC(serialize_getint IRCALL_lj_serialize_getint)
{
	serialize_context* ctx = serialize_toctx(L);
	setintV(L->top++, lj_serialize_getint(ctx));
	return 1;
}

LJLIB_CF(serialize_nextmsg)  LJLIB_REC(serialize_nextmsg)
{
	serialize_context* ctx = serialize_toctx(L);
	setintV(L->top++, lj_serialize_nextmsg(ctx));
	return 1;
}

LJLIB_CF(serialize_firstmsg)  LJLIB_REC(serialize_firstmsg)
{
	serialize_context* ctx = serialize_toctx(L);
	setintV(L->top++, lj_serialize_firstmsg(ctx));
	return 1;
}

LJLIB_CF(serialize_print_state)
{
	serialize_context* ctx = serialize_toctx(L);

	printf("C: [%x, %x)\n", ctx->c_low, ctx->c_low + ctx->c_range);
	printf("D: [%x, %x, %x)\n", ctx->d_low, ctx->d_code, ctx->d_low + ctx->d_range);
	printf("Read bytes left: %d\n", ctx->reader.end - ctx->reader.next);
	return 0;
}

/* ------------------------------------------------------------------------ */

#include "../lj_libdef.h"

LUALIB_API int luaopen_serialize(lua_State *L)
{
	LJ_LIB_REG_(L, NULL, serialize_method);
	copyTV(L, L->top, L->top-1); L->top++;
	lua_setfield(L, LUA_REGISTRYINDEX, LUA_SERIALIZE_CONTEXT);
	LJ_LIB_REG(L, serialize);
	return 1;
}

